import torch
from system_initializer import SystemInitializer
from local_training import LocalTraining
from gradient_predictor import GradientPredictor
from adaptive_noise import AdaptiveNoiseInjector
from dynamic_aggregator import DynamicWeightAggregator
from privacy_budget import PrivacyBudgetManager
from models import LeNet5
from dataset_loader import load_datasets, get_data_loaders

def main():
    # 实验参数 (对应论文实验设置)
    num_clients = 20
    global_rounds = 100
    base_iterations = 5
    total_privacy_budget = 10.0
    dataset_name = "MNIST"  # 或 "CIFAR-10"
    
    # 初始化系统
    initializer = SystemInitializer(num_clients, total_privacy_budget, base_iterations)
    
    # 加载数据集
    client_datasets, val_dataset, testset = load_datasets(dataset_name, num_clients)
    
    # 初始化服务器
    num_channels = 1 if dataset_name == "MNIST" else 3
    server_state = initializer.server_initialization(LeNet5, num_channels=num_channels)
    
    # 初始化客户端
    clients = []
    for i in range(num_clients):
        client_state = initializer.client_initialization(i, client_datasets[i], LeNet5, num_channels=num_channels)
        clients.append(client_state)
    
    # 初始化各模块
    gradient_predictors = [GradientPredictor() for _ in range(num_clients)]
    noise_injector = AdaptiveNoiseInjector()
    weight_aggregator = DynamicWeightAggregator()
    
    # 初始化隐私预算管理器 (假设敏感度等级随机生成)
    sensitivity_levels = np.random.rand(num_clients)
    privacy_manager = PrivacyBudgetManager(total_privacy_budget, sensitivity_levels)
    
    # 训练循环
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    for round in range(global_rounds):
        # 训练步骤 (此处省略具体实现，参考之前完整代码)
        # ...
        
        # 每10轮评估模型性能
        if round % 10 == 0:
            accuracy = evaluate_model(server_state['global_model'], testset, device)
            remaining_budget = privacy_manager.get_remaining_budget()
            print(f"轮次 {round}: 测试准确率 = {accuracy:.2f}%, 剩余预算 = {remaining_budget:.2f}")

def evaluate_model(model, testset, device, batch_size=64):
    """评估模型性能"""
    test_loader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False)
    model.eval()
    model.to(device)
    
    correct = 0
    total = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            outputs = model(data)
            _, predicted = torch.max(outputs.data, 1)
            total += target.size(0)
            correct += (predicted == target).sum().item()
    
    return 100 * correct / total

if __name__ == "__main__":
    main()
